import type { OpenAPIObject } from '@nestjs/swagger';
import deepmerge from 'deepmerge';
import { JSONSchema } from 'zod/v4/core';
import { fixAllRefs, convertToOpenApi3Point0 } from './utils';
import { DEFS_KEY, EMPTY_TYPE_KEY, HAS_CONST_KEY, HAS_NULL_KEY, PARENT_ADDITIONAL_PROPERTIES_KEY, PARENT_HAS_REFS_KEY, PARENT_ID_KEY, UNWRAP_ROOT_KEY } from './const';
import { isDeepStrictEqual } from 'node:util';
import { assert } from './assert';

type DtoSchema = Exclude<Exclude<OpenAPIObject['components'], undefined>['schemas'], undefined>[string];
type OpenAPIParameter = Exclude<Exclude<Exclude<OpenAPIObject['paths'], undefined>['/'], undefined>['parameters'], undefined>[number];

/**
 * This function performs some post-processing on the OpenAPI document.  It
 * should only touch parts of the document that were generated from nestjs-zod
 * DTOs.
 * 
 * Specifically, this function:
 * 1. Removes empty `type` fields
 * 2. Renames OpenAPI schemas that have an explicit `id` field to match that
 *    `id`, instead of using the DTO class name
 * 3. If the DTO's schema references another zod schema, it adds that zod
 *    schema's OpenAPI representation to `components.schemas`
 * 4. If a DTO is created directly with an array zod schema, it ensures the
 *    OpenAPI schema is generated properly
 * 5. Handles recursive zod schemas
 * 6. Handles `null` properly based on the OpenAPI version
 * 
 * @param doc - The OpenAPI document that is generated by `SwaggerModule.createDocument`
 * @param options.version - The version of OpenAPI to use.  Defaults to `auto`,
 * which will use the version of the OpenAPI object passed in.  Note if the
 * version is `3.1` then it impacts how `null` is handled.  In `3.0`, `nullable:
 * true` is used, while in `3.1` `anyOf: [..., { type: 'null' }]` is used
 * instead
 * @returns A cleaned up OpenAPI document
 */
export function cleanupOpenApiDoc(doc: OpenAPIObject, { version: versionParam = 'auto' }: { version?: '3.1' | '3.0' | 'auto' } = {}): OpenAPIObject {
    const schemas: Record<string, DtoSchema> = {};
    const renames: Record<string, string> = {};
    const version = versionParam === 'auto' ? (doc.openapi.startsWith('3.1') ? '3.1' : '3.0') : versionParam;

    for (let [oldSchemaName, oldOpenapiSchema] of Object.entries(doc.components?.schemas || {})) {
        // Ignore non-object types, which are not added by us
        if (!('type' in oldOpenapiSchema) || oldOpenapiSchema.type !== 'object') {
            schemas[oldSchemaName] = oldOpenapiSchema;
            continue;
        }

        let newSchemaName = oldSchemaName;
        let addedDefs = false;
        let hasRefs = false;
        let hasNull = false;
        let hasConst = false;
        const defRenames: Record<string, string> = {};

        // Clone so we can mutate
        let newOpenapiSchema = deepmerge<typeof oldOpenapiSchema>({}, oldOpenapiSchema);

        for (let propertySchema of Object.values(newOpenapiSchema.properties || {})) {
            if (HAS_CONST_KEY in propertySchema) {
                hasConst = Boolean(propertySchema[HAS_CONST_KEY]);
                delete propertySchema[HAS_CONST_KEY];
            }

            if (HAS_NULL_KEY in propertySchema) {
                hasNull = Boolean(propertySchema[HAS_NULL_KEY]);
                delete propertySchema[HAS_NULL_KEY];
            }

            if (PARENT_HAS_REFS_KEY in propertySchema) {
                hasRefs = Boolean(propertySchema[PARENT_HAS_REFS_KEY]);
                delete propertySchema[PARENT_HAS_REFS_KEY];
            }

            // Remove `type` if we added `type: ''`
            if (EMPTY_TYPE_KEY in propertySchema && propertySchema[EMPTY_TYPE_KEY]) {
                delete propertySchema[EMPTY_TYPE_KEY];
                if ('type' in propertySchema && propertySchema.type === '') {
                    delete propertySchema.type;
                }
            }

            // Rename the schema if using `meta({ id: "NewName" })`
            if (PARENT_ID_KEY in propertySchema && typeof propertySchema[PARENT_ID_KEY] === 'string') {
                Object.assign(newOpenapiSchema, { id: propertySchema[PARENT_ID_KEY] });
                newSchemaName = propertySchema[PARENT_ID_KEY];
                delete propertySchema[PARENT_ID_KEY];
            }

            if (PARENT_ADDITIONAL_PROPERTIES_KEY in propertySchema && typeof propertySchema[PARENT_ADDITIONAL_PROPERTIES_KEY] === 'boolean') {
                newOpenapiSchema.additionalProperties = propertySchema[PARENT_ADDITIONAL_PROPERTIES_KEY];
                delete propertySchema[PARENT_ADDITIONAL_PROPERTIES_KEY];
            }

            // Add each $def as a schema
            if (DEFS_KEY in propertySchema) {
                const defs = propertySchema[DEFS_KEY] as Record<string, JSONSchema.BaseSchema>;
                delete propertySchema[DEFS_KEY];

                if (!addedDefs) {
                    // If the def has no ID, then we need to prefix the def key
                    // with the root schema name to make it globally unique 
                    // This can happen if the def is part of a recursive schema
                    // (for example, `___schema0`)
                    for (let [defSchemaId, defSchema] of Object.entries(defs)) {
                        if (!('id' in defSchema)) {
                            defRenames[defSchemaId] = `${newSchemaName}${defSchemaId}`;
                        }
                    }

                    for (let [defSchemaId, defSchema] of Object.entries(defs)) {
                        let fixedDef = fixAllRefs({ schema: defSchema, rootSchemaName: newSchemaName, defRenames })  

                        if (version === '3.0') {
                            fixedDef = convertToOpenApi3Point0(fixedDef);
                        }

                        const newDefSchemaKey = defRenames[defSchemaId] || defSchemaId;

                        if (schemas[newDefSchemaKey] && !isDeepStrictEqual(schemas[newDefSchemaKey], fixedDef)) {
                            throw new Error(`[cleanupOpenApiDoc] Found multiple schemas with name \`${newDefSchemaKey}\`.  Please review your schemas to ensure that you are not using the same schema name for different schemas`);
                        }

                        // @ts-ignore TODO: fix this
                        schemas[newDefSchemaKey] = fixedDef;
                    }

                    addedDefs = true;
                }
            }
        }

        if (newSchemaName !== oldSchemaName) {
            renames[oldSchemaName] = newSchemaName;

            // @ts-expect-error TODO: is ID a valid openapi field?
            newOpenapiSchema['id'] = newSchemaName;
        }

        if (hasRefs) {
            // @ts-ignore TODO: fix this
            newOpenapiSchema = fixAllRefs({
                // @ts-expect-error TODO: fix TS error
                schema: newOpenapiSchema,
                rootSchemaName: newSchemaName,
                defRenames,
            });
        }

        // Zod generates openapi schemas like this when a field is nullable:
        //
        // {
        //   anyOf: [
        //     { type: 'string' },
        //     { type: 'null' }
        //   ]
        // }
        //
        // However, this is not valid openapi in 3.0.  So if the user wants to generate openapi 3.0 docs,
        // we convert it to this:
        //
        // { 
        //   type: 'string', 
        //   nullable: true 
        // }
        //
        // This is the default behavior, since nestjs/swagger seems to generate
        // a 3.0 document by default
        if ((hasNull || hasConst) && version === '3.0') {
            // @ts-expect-error TODO: fix this
            newOpenapiSchema = convertToOpenApi3Point0(newOpenapiSchema);
        }

        // When the consumer does this `createZodDto(z.array(...))`, then
        // `_OPENAPI_METADATA_FACTORY` will return: 
        // `{ root: { type: 'array', [UNWRAP_ROOT_KEY]: true } }`
        // This is a workaround for the fact that the factory doesn't support
        // returning array schemas directly.  
        // If we see `UNWRAP_ROOT_KEY`, then we unwrap the array
        // schema (replace the root object schema with the array schema)
        if (newOpenapiSchema.properties && 'root' in newOpenapiSchema.properties && UNWRAP_ROOT_KEY in newOpenapiSchema.properties.root) {
            const replaceRoot = newOpenapiSchema.properties.root[UNWRAP_ROOT_KEY];
            delete newOpenapiSchema.properties.root[UNWRAP_ROOT_KEY];
            if (replaceRoot) {
                // @ts-expect-error TODO: fix this
                newOpenapiSchema = newOpenapiSchema.properties.root;
            }
        }

        if (schemas[newSchemaName] && !isDeepStrictEqual(schemas[newSchemaName], newOpenapiSchema)) {
            throw new Error(`[cleanupOpenApiDoc] Found multiple schemas with name \`${newSchemaName}\`.  Please review your schemas to ensure that you are not using the same schema name for different schemas`);
        }
        
        schemas[newSchemaName] = newOpenapiSchema;
    }

    // Rename all the references for 
    const paths = deepmerge<typeof doc.paths>(doc.paths, {})
    for (let { get, patch, post, delete: del, put, head } of Object.values(paths)) {
        for (let methodObject of Object.values({ get, patch, post, del, put, head })) {
            const content = methodObject?.requestBody && 'content' in methodObject?.requestBody && methodObject?.requestBody.content || {}
            for (let requestBodyObject of Object.values(content)) {
                if (requestBodyObject.schema && '$ref' in requestBodyObject.schema) {
                    const oldSchemaName = getSchemaNameFromRef(requestBodyObject.schema.$ref);
                    if (renames[oldSchemaName]) {
                        const newSchemaName = renames[oldSchemaName];
                        requestBodyObject.schema.$ref = requestBodyObject.schema.$ref.replace(`/${oldSchemaName}`, `/${newSchemaName}`);
                    }
                }
            }

            for (let statusCodeObject of Object.values(methodObject?.responses || {})) {
                const content = statusCodeObject && 'content' in statusCodeObject && statusCodeObject.content || {};
                for (let responseBodyObject of Object.values(content)) {
                    if (responseBodyObject.schema && '$ref' in responseBodyObject.schema) {
                        const oldSchemaName = getSchemaNameFromRef(responseBodyObject.schema.$ref);
                        if (renames[oldSchemaName]) {
                            const newSchemaName = renames[oldSchemaName];
                            responseBodyObject.schema.$ref = responseBodyObject.schema.$ref.replace(`/${oldSchemaName}`, `/${newSchemaName}`);
                        }
                    }
                }
            }

            for (let i = 0; i < (methodObject?.parameters || []).length; i++) {
                assert(methodObject?.parameters, 'parameters is required');

                let parameter = methodObject.parameters[i];
                parameter = fixParameter(parameter, version);

                // Add each $def as a schema
                if (DEFS_KEY in parameter) {
                    const defs = parameter[DEFS_KEY] as Record<string, JSONSchema.BaseSchema>;
                    delete parameter[DEFS_KEY];

                    for (let [defSchemaId, defSchema] of Object.entries(defs)) {                        
                        let fixedDef;
                        try {
                            fixedDef = fixAllRefs({ schema: defSchema });
                        } catch (err) {
                            if (err instanceof Error && err.message.startsWith('[fixAllRefs]')) {
                                throw new Error(`[cleanupOpenApiDoc] Recursive schemas are not supported for parameters`, { cause: err });
                            }
                            throw err;
                        }

                        if (schemas[defSchemaId] && !isDeepStrictEqual(schemas[defSchemaId], fixedDef)) {
                            throw new Error(`[cleanupOpenApiDoc] Found multiple schemas with name \`${defSchemaId}\`.  Please review your schemas to ensure that you are not using the same schema name for different schemas`);
                        }

                        // @ts-ignore TODO: fix this
                        schemas[defSchemaId] = fixedDef;
                    }
                }

                methodObject.parameters[i] = parameter;
            }
        }
    }

    return {
        ...doc,
        paths,
        components: {
            ...doc.components,
            schemas,
        }
    }
}

/**
 * Fixes various issues with parameters:
 * 1. Removes empty `type` fields
 * 2. Moves `const` from the root level of the parameter object to the schema
 * 3. Fixes refs to point to components.schemas
 * 4. Removes some nestjs-zod markers
 */
function fixParameter(parameterInput: OpenAPIParameter, version: '3.1' | '3.0') {
    const parameter = deepmerge<typeof parameterInput>({}, parameterInput);

    // nestjs seems to move some stuff out of the schema and into the root level of the parameter object 🤷
    if (EMPTY_TYPE_KEY in parameter) {
        delete parameter[EMPTY_TYPE_KEY];
        if ('schema' in parameter && parameter.schema && 'type' in parameter.schema) {
            delete parameter.schema.type;
        }
    }

    if (UNWRAP_ROOT_KEY in parameter) {
        throw new Error(`[cleanupOpenApiDoc] Query or url parameters must be an object type`);
    }

    if ('const' in parameter && 'schema' in parameter && parameter.schema) {
        Object.assign(parameter.schema, { const: parameter.const });
        delete parameter.const;
    }

    if (PARENT_ID_KEY in parameter) {
        delete parameter[PARENT_ID_KEY];
    }

    if (PARENT_HAS_REFS_KEY in parameter) {
        delete parameter[PARENT_HAS_REFS_KEY];

        if ('schema' in parameter) {
            try {
                // @ts-expect-error TODO: fix this
                parameter.schema = fixAllRefs({ schema: parameter.schema });
            } catch (err) {
                if (err instanceof Error && err.message.startsWith('[fixAllRefs]')) {
                    throw new Error(`[cleanupOpenApiDoc] Recursive schemas are not supported for parameters`, { cause: err });
                }
                throw err;
            }
        }
    }

    if (HAS_CONST_KEY in parameter) {
        delete parameter[HAS_CONST_KEY];

        if (version === '3.0' && 'schema' in parameter && parameter.schema) {
            // @ts-expect-error TODO: fix this
            parameter.schema = convertToOpenApi3Point0(parameter.schema);
        }
    }

    return parameter;
}

function getSchemaNameFromRef(ref: string) {
    const lastSlash = ref.lastIndexOf("/");
    const schemaName = ref.slice(lastSlash + 1);
    return schemaName;
}
